"""代理图定义模块 - 包含所有代理的状态图定义。"""

import sys
import os
import sqlite3
from typing import Literal
from langgraph.graph import END, START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langchain_core.messages import AIMessage, trim_messages
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from concurrent.futures import ThreadPoolExecutor

# # 添加项目根目录到 Python 路径
# project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
# src_path = os.path.join(project_root, "src")
# if src_path not in sys.path:
#     sys.path.insert(0, src_path)

# 现在可以使用绝对导入
from agents.configurations import ReactConfiguration, SqlConfiguration
from agents.states import (
    ReactGraphAnnotation,
    SqlGraphAnnotation,
    SupervisorState,
    DocumentGraphAnnotation,
)
from nodes.react_nodes import generate_response
from nodes.tool_nodes import react_tools_node, get_schema_node, run_query_node
from nodes.sql_nodes import list_tables, call_get_schema, generate_query, check_query, trim_memory
from nodes.document_nodes import document_reader_node, document_analyzer_node
from utils.models import load_chat_model
from utils.logger import logger
from utils.database_pool import (
    get_async_connection,
    reset_connection_pool,
    monitor_connection_pool,
    ConnectionPoolError,
)


# 自定义函数修剪和过滤state中的消息
def filter_messages(messages: list):
    if len(messages) <= 3:
        return messages
    return messages[-3:]


def should_continue(state: SqlGraphAnnotation) -> Literal[END, "check_query"]:
    messages = state["messages"]
    last_message = messages[-1]
    if not last_message.tool_calls:
        state["messages"] = filter_messages(state["messages"])
        return END
    else:
        return "check_query"


# React 代理图定义
def create_react_graph():
    """创建 React 代理的状态图。"""
    builder = StateGraph(ReactGraphAnnotation, config_schema=ReactConfiguration)

    builder.add_node("generate_response", generate_response)
    builder.add_node("tools", react_tools_node)

    builder.add_edge(START, "generate_response")
    builder.add_conditional_edges("generate_response", tools_condition, ["tools", END])
    builder.add_edge("tools", "generate_response")

    graph = builder.compile()
    graph.name = "react_agent"
    return graph


# SQL 代理图定义
def create_sql_graph():
    """创建 SQL 代理的状态图（同步版本）。"""
    try:
        # 尝试使用SQLite检查点保存器
        db_path = "checkpoints.db"
        os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)

        # 测试SQLite连接
        with sqlite3.connect(db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
            result = cursor.fetchone()
            if result != (1,):
                raise ConnectionPoolError("SQLite连接测试失败")

        memory = SqliteSaver.from_conn_string(db_path)
        logger.info("Using SQLite checkpointer for SQL graph")
    except Exception as e:
        logger.warning(
            f"Failed to initialize SQLite checkpointer: {e}, falling back to InMemorySaver"
        )
        memory = InMemorySaver()

    builder = StateGraph(SqlGraphAnnotation, config_schema=SqlConfiguration)

    # 添加节点
    builder.add_node(list_tables)
    builder.add_node(call_get_schema)
    builder.add_node("get_schema", get_schema_node)
    builder.add_node(generate_query)
    builder.add_node(check_query)
    builder.add_node("run_query", run_query_node)
    builder.add_node(trim_memory)
    # 添加边
    builder.add_edge(START, "list_tables")
    builder.add_edge("list_tables", "call_get_schema")
    builder.add_edge("call_get_schema", "get_schema")
    builder.add_edge("get_schema", "generate_query")

    # 同时需要更新条件边的定义
    builder.add_conditional_edges(
        "generate_query",
        should_continue,
    )

    builder.add_edge("check_query", "run_query")
    builder.add_edge("run_query", "generate_query")

    graph = builder.compile(checkpointer=memory)
    graph.name = "sql_agent"
    return graph


async def create_sql_graph_async():
    """创建 SQL 代理的状态图（异步版本，使用连接池管理）。"""
    try:
        # 重置连接池状态以避免线程重复启动问题
        await reset_connection_pool()

        # 使用连接池获取连接
        conn = await get_async_connection()
        builder = StateGraph(SqlGraphAnnotation, config_schema=SqlConfiguration)

        # 创建 AsyncSqliteSaver 实例
        saver = AsyncSqliteSaver(conn)

        # 添加节点
        builder.add_node(list_tables)
        builder.add_node(call_get_schema)
        builder.add_node("get_schema", get_schema_node)
        builder.add_node(generate_query)
        builder.add_node(check_query)
        builder.add_node("run_query", run_query_node)

        # 添加边
        builder.add_edge(START, "list_tables")
        builder.add_edge("list_tables", "call_get_schema")
        builder.add_edge("call_get_schema", "get_schema")
        builder.add_edge("get_schema", "generate_query")

        # 同时需要更新条件边的定义
        builder.add_conditional_edges(
            "generate_query",
            should_continue,
        )

        builder.add_edge("check_query", "run_query")
        builder.add_edge("run_query", "generate_query")

        graph = builder.compile(checkpointer=saver)
        graph.name = "sql_agent"
        logger.info("SQL Agent graph created successfully with AsyncSqliteSaver")
        return graph

    except Exception as e:
        logger.error(f"Failed to create async SQL graph: {e}")
        logger.info("Falling back to synchronous SQL graph")
        return create_sql_graph()


# 文档解析图定义
from IPython.display import display, Image
import os

# 创建图实例
react_graph = create_react_graph()

# 启动连接池监控
monitor_connection_pool()

os.makedirs("output", exist_ok=True)

# 创建SQL图实例
sql_graph = create_sql_graph()
